package com.yricky.atri.model

import com.google.gson.Gson
import com.yricky.atri.globalContext
import com.yricky.atri.model.pojo.ModelMetaData
import com.yricky.atri.model.pojo.TensorMetaData
import org.tensorflow.lite.Interpreter
import org.tensorflow.lite.Tensor
import java.io.File
import java.io.InputStream
import kotlin.collections.ArrayList

/**
 * @author Yricky
 * @date 2021/12/8
 */
object MLModelManager {
    private val dir: File by lazy {
        globalContext.getExternalFilesDir("models")!!
    }

    fun importModel(name:String,model:InputStream):Boolean{
        if(name.contains("/"))
            return false
        val modelDir = File(dir,name)
        if(modelDir.exists())
            return false
        modelDir.mkdirs()
        val modelFile = File(modelDir,"model")
        modelFile.writeBytes(model.readBytes())
        return try {
            val metaData:ModelMetaData = Interpreter(modelFile).use { interpreter->
                val input = ArrayList<TensorMetaData>()
                val output = ArrayList<TensorMetaData>()
                (0 until interpreter.inputTensorCount).forEach {
                    input.add(interpreter.getInputTensor(it).getMetaData())
                }
                (0 until interpreter.outputTensorCount).forEach {
                    output.add(interpreter.getOutputTensor(it).getMetaData())
                }
                ModelMetaData(name,input, output)
            }
            File(modelDir,"metaData.json").writeText(Gson().toJson(metaData))
            true
        }catch (e:Throwable){
            modelDir.deleteRecursively()
            false
        }
    }

    fun getModels():List<Pair<File,ModelMetaData>>{
        return dir.listFiles { f->f.isDirectory }?.flatMap {
            try {
                listOf(
                    File(it,"model") to (Gson().fromJson(File(it,"metaData.json").readText(),ModelMetaData::class.java))
                )
            }catch (e:Throwable){
                emptyList()
            }
        }?: emptyList()
    }


    private fun Tensor.getMetaData():TensorMetaData{
        return TensorMetaData(
            name=name(),
            dataType = dataType(),
            numDimensions = numDimensions(),
            shape = shape(),
            numBytes = numBytes(),
            numElements = numElements()
        )
    }
}